#include <jni.h>
#include <string>
#include <memory>
#include "include/score_engine.h"
#include "log.h"
// 常量定义
#define ENGINE_VERSION "1.0.0"

// JNI 命名宏，简化函数命名
#define JNI_METHOD(return_type, method_name) \
    extern "C" JNIEXPORT return_type JNICALL \
    Java_com_thunder_ktv_scoreengine2sdk_ScoreEngine2_##method_name


typedef struct {
    uint32_t start_frame;   //字起始帧索引
    uint32_t end_frame;     //字结束帧索引
    float energy;           //区间平均能量
    float pitch;            //区间音调中位数
    char text[16];          //字符串，不足补 0
}CharEntry; // 32 Bytes


/**
 * 将Java的ScoreEngineConfig转换为C++的ScoreEngineConfig
 * 
 * @param env JNI环境
 * @param jConfig Java配置对象
 * @return C++配置对象
 */
ScoreEngineConfig convertJavaConfig(JNIEnv* env, jobject jConfig) {
    ScoreEngineConfig config;
    memset(&config, 0, sizeof(ScoreEngineConfig));
    
    if (jConfig == nullptr) {
        // 使用默认配置
        config.enabled_dimensions = SCORE_ALL;
        config.style = SONG_STYLE_POP;
        config.difficulty = DIFFICULTY_NORMAL;
        config.custom_weights = nullptr;
        config.custom_factors = nullptr;
        return config;
    }
    
    // 获取Java类引用
    jclass configClass = env->GetObjectClass(jConfig);
    
    // 获取字段ID
    jfieldID enabledDimensionsField = env->GetFieldID(configClass, "enabledDimensions", "I");
    jfieldID styleField = env->GetFieldID(configClass, "style", "Lcom/thunder/ktv/scoreengine2sdk/SongStyle;");
    jfieldID difficultyField = env->GetFieldID(configClass, "difficulty", "Lcom/thunder/ktv/scoreengine2sdk/DifficultyLevel;");
    jfieldID weightsField = env->GetFieldID(configClass, "weights", "Lcom/thunder/ktv/scoreengine2sdk/ScoreEngine2$ScoreWeights;");
    jfieldID factorsField = env->GetFieldID(configClass, "factors", "Lcom/thunder/ktv/scoreengine2sdk/ScoreEngine2$DifficultyFactors;");
    
    // 获取基本字段值
    config.enabled_dimensions = env->GetIntField(jConfig, enabledDimensionsField);
    
    // 获取SongStyle枚举值
    jobject jStyleEnum = env->GetObjectField(jConfig, styleField);
    if (jStyleEnum != nullptr) {
        jclass styleEnumClass = env->GetObjectClass(jStyleEnum);
        jmethodID getValueMethod = env->GetMethodID(styleEnumClass, "getValue", "()I");
        int style = env->CallIntMethod(jStyleEnum, getValueMethod);
        
        switch (style) {
            case 0: config.style = SONG_STYLE_POP; break;
            case 1: config.style = SONG_STYLE_ROCK; break;
            case 2: config.style = SONG_STYLE_FOLK; break;
            case 3: config.style = SONG_STYLE_RAP; break;
            case 4: config.style = SONG_STYLE_RNB; break;
            case 5: config.style = SONG_STYLE_CLASSICAL; break;
            case 6: config.style = SONG_STYLE_JAZZ; break;
            case 7: config.style = SONG_STYLE_ELECTRONIC; break;
            case 8: config.style = SONG_STYLE_CUSTOM; break;
            default: config.style = SONG_STYLE_POP; break;
        }
    } else {
        config.style = SONG_STYLE_POP; // 默认
    }
    
    // 获取DifficultyLevel枚举值
    jobject jDifficultyEnum = env->GetObjectField(jConfig, difficultyField);
    if (jDifficultyEnum != nullptr) {
        jclass difficultyEnumClass = env->GetObjectClass(jDifficultyEnum);
        jmethodID getValueMethod = env->GetMethodID(difficultyEnumClass, "getValue", "()I");
        int difficulty = env->CallIntMethod(jDifficultyEnum, getValueMethod);
        
        switch (difficulty) {
            case 0: config.difficulty = DIFFICULTY_EASY; break;
            case 1: config.difficulty = DIFFICULTY_NORMAL; break;
            case 2: config.difficulty = DIFFICULTY_HARD; break;
            case 3: config.difficulty = DIFFICULTY_EXPERT; break;
            case 4: config.difficulty = DIFFICULTY_CUSTOM; break;
            default: config.difficulty = DIFFICULTY_NORMAL; break;
        }
    } else {
        config.difficulty = DIFFICULTY_NORMAL; // 默认
    }
    
    // 获取自定义权重对象
    jobject jWeights = env->GetObjectField(jConfig, weightsField);
    if (jWeights != nullptr) {
        // 创建C++权重结构体
        ScoreWeights* weights = new ScoreWeights();
        memset(weights, 0, sizeof(ScoreWeights));
        
        // 获取权重类引用
        jclass weightsClass = env->GetObjectClass(jWeights);
        
        // 获取字段ID
        jfieldID pitchWeightField = env->GetFieldID(weightsClass, "pitchWeight", "F");
        jfieldID brightnessWeightField = env->GetFieldID(weightsClass, "brightnessWeight", "F");
        jfieldID rhythmWeightField = env->GetFieldID(weightsClass, "rhythmWeight", "F");
        jfieldID emotionWeightField = env->GetFieldID(weightsClass, "emotionWeight", "F");
        jfieldID breathWeightField = env->GetFieldID(weightsClass, "breathWeight", "F");
        jfieldID timbreWeightField = env->GetFieldID(weightsClass, "timbreWeight", "F");
        
        // 获取权重值
        weights->pitch_weight = env->GetFloatField(jWeights, pitchWeightField);
        weights->brightness_weight = env->GetFloatField(jWeights, brightnessWeightField);
        weights->rhythm_weight = env->GetFloatField(jWeights, rhythmWeightField);
        weights->emotion_weight = env->GetFloatField(jWeights, emotionWeightField);
        weights->breath_weight = env->GetFloatField(jWeights, breathWeightField);
        weights->timbre_weight = env->GetFloatField(jWeights, timbreWeightField);
        
        // 设置到配置对象
        config.custom_weights = weights;
    } else {
        config.custom_weights = nullptr;
    }
    
    // 获取自定义难度因子对象
    jobject jFactors = env->GetObjectField(jConfig, factorsField);
    if (jFactors != nullptr) {
        // 创建C++难度因子结构体
        DifficultyFactors* factors = new DifficultyFactors();
        memset(factors, 0, sizeof(DifficultyFactors));
        
        // 获取难度因子类引用
        jclass factorsClass = env->GetObjectClass(jFactors);
        
        // 获取字段ID
        jfieldID pitchToleranceField = env->GetFieldID(factorsClass, "pitchTolerance", "F");
        jfieldID rhythmToleranceField = env->GetFieldID(factorsClass, "rhythmTolerance", "F");
        jfieldID otherTolerancesField = env->GetFieldID(factorsClass, "otherTolerances", "F");
        
        // 获取难度因子值
        factors->pitch_tolerance = env->GetFloatField(jFactors, pitchToleranceField);
        factors->rhythm_tolerance = env->GetFloatField(jFactors, rhythmToleranceField);
        factors->other_tolerances = env->GetFloatField(jFactors, otherTolerancesField);
        
        // 设置到配置对象
        config.custom_factors = factors;
    } else {
        config.custom_factors = nullptr;
    }
    
    return config;
}

/**
 * 将C++的ScoreResult填充到Java的ScoreResult
 * 
 * @param env JNI环境
 * @param cResult C++结果对象
 * @param jResult Java结果对象
 */
void fillJavaResult(JNIEnv* env, const ScoreResult& cResult, jobject jResult) {
    // 获取Java类引用
    jclass resultClass = env->GetObjectClass(jResult);
    
    // 获取字段ID
    jfieldID scoreFlagField = env->GetFieldID(resultClass, "scoreFlag", "I");
    jfieldID overallScoreField = env->GetFieldID(resultClass, "overallScore", "I");
    jfieldID pitchScoreField = env->GetFieldID(resultClass, "pitchScore", "I");
    jfieldID brightnessScoreField = env->GetFieldID(resultClass, "brightnessScore", "I");
    jfieldID rhythmScoreField = env->GetFieldID(resultClass, "rhythmScore", "I");
    jfieldID emotionScoreField = env->GetFieldID(resultClass, "emotionScore", "I");
    jfieldID breathScoreField = env->GetFieldID(resultClass, "breathScore", "I");
    jfieldID timbreScoreField = env->GetFieldID(resultClass, "timbreScore", "I");
    jfieldID totalFramesField = env->GetFieldID(resultClass, "totalFrames", "I");
    jfieldID frameFeatureField = env->GetFieldID(resultClass, "frameFeature", "Lcom/thunder/ktv/scoreengine2sdk/ScoreEngine2$FrameFeature;");
    jfieldID baseFrameFeatureField = env->GetFieldID(resultClass, "baseFrameFeature", "Lcom/thunder/ktv/scoreengine2sdk/ScoreEngine2$FrameFeature;");

    // 设置基本字段值
    env->SetIntField(jResult, scoreFlagField, cResult.score_flag);
    env->SetIntField(jResult, overallScoreField, cResult.overall_score);
    env->SetIntField(jResult, pitchScoreField, cResult.pitch_score);
    env->SetIntField(jResult, brightnessScoreField, cResult.brightness_score);
    env->SetIntField(jResult, rhythmScoreField, cResult.rhythm_score);
    env->SetIntField(jResult, emotionScoreField, cResult.emotion_score);
    env->SetIntField(jResult, breathScoreField, cResult.breath_score);
    env->SetIntField(jResult, timbreScoreField, cResult.timbre_score);
    env->SetIntField(jResult, totalFramesField, cResult.total_frames);
    
    // 获取frameFeature对象
    jobject jFrameFeature = env->GetObjectField(jResult, frameFeatureField);
    if (jFrameFeature != nullptr) {
        // 获取FrameFeature类引用
        jclass frameFeatureClass = env->GetObjectClass(jFrameFeature);
        
        // 获取字段ID
        jfieldID frameIndexField = env->GetFieldID(frameFeatureClass, "frameIndex", "I");
        jfieldID vadValueField = env->GetFieldID(frameFeatureClass, "vadValue", "F");
        jfieldID energyField = env->GetFieldID(frameFeatureClass, "energy", "F");
        jfieldID pitchField = env->GetFieldID(frameFeatureClass, "pitch", "F");
        jfieldID brightnessField = env->GetFieldID(frameFeatureClass, "brightness", "F");
        jfieldID timbreField = env->GetFieldID(frameFeatureClass, "timbre", "F");
        
        // 设置帧特征值
        env->SetIntField(jFrameFeature, frameIndexField, cResult.frame_feature.frame_index);
        env->SetFloatField(jFrameFeature, vadValueField, cResult.frame_feature.vad_value);
        env->SetFloatField(jFrameFeature, energyField, cResult.frame_feature.energy);
        env->SetFloatField(jFrameFeature, pitchField, cResult.frame_feature.pitch);
        env->SetFloatField(jFrameFeature, brightnessField, cResult.frame_feature.brightness);
        env->SetFloatField(jFrameFeature, timbreField, cResult.frame_feature.timbre);
    }
    // 获取baseFrameFeature对象
    jobject jBaseFrameFeature = env->GetObjectField(jResult, baseFrameFeatureField);
    if (jBaseFrameFeature != nullptr) {
        // 获取FrameFeature类引用
        jclass baseFrameFeatureClass = env->GetObjectClass(jBaseFrameFeature);

        // 获取字段ID
        jfieldID frameIndexField = env->GetFieldID(baseFrameFeatureClass, "frameIndex", "I");
        jfieldID vadValueField = env->GetFieldID(baseFrameFeatureClass, "vadValue", "F");
        jfieldID energyField = env->GetFieldID(baseFrameFeatureClass, "energy", "F");
        jfieldID pitchField = env->GetFieldID(baseFrameFeatureClass, "pitch", "F");
        jfieldID brightnessField = env->GetFieldID(baseFrameFeatureClass, "brightness", "F");
        jfieldID timbreField = env->GetFieldID(baseFrameFeatureClass, "timbre", "F");

        // 设置帧特征值
        env->SetIntField(jBaseFrameFeature, frameIndexField, cResult.base_frame_feature.frame_index);
        env->SetFloatField(jBaseFrameFeature, vadValueField, cResult.base_frame_feature.vad_value);
        env->SetFloatField(jBaseFrameFeature, energyField, cResult.base_frame_feature.energy);
        env->SetFloatField(jBaseFrameFeature, pitchField, cResult.base_frame_feature.pitch);
        env->SetFloatField(jBaseFrameFeature, brightnessField, cResult.base_frame_feature.brightness);
        env->SetFloatField(jBaseFrameFeature, timbreField, cResult.base_frame_feature.timbre);

        // LOGD("cResult.base_frame_feature.frame_index %d\n",cResult.base_frame_feature.frame_index);
        // LOGD("cResult.base_frame_feature.vad_value %f\n",cResult.base_frame_feature.vad_value);
        // LOGD("cResult.base_frame_feature.pitch %f\n",cResult.base_frame_feature.pitch);
        // LOGD("cResult.base_frame_feature.energy %f\n",cResult.base_frame_feature.energy);
        // LOGD("cResult.base_frame_feature.brightness %f\n",cResult.base_frame_feature.brightness);
        // LOGD("cResult.base_frame_feature.timbre %f\n",cResult.base_frame_feature.timbre);
    }
}

/**
 * 初始化评分引擎
 */
JNI_METHOD(jboolean, initScoreEngine)(JNIEnv* env, jobject thiz, jbyteArray jBaseFileData, 
                                     jint jDataSize, jint jFrameSize, jint jSampleRate, jobject jConfig) {
    // 获取Java类引用
    jclass clazz = env->GetObjectClass(thiz);
    LOGD("native ScoreEngineVersion v1.5\n");
    // 获取引擎句柄字段ID
    jfieldID handleField = env->GetFieldID(clazz, "nativeEngineHandle", "J");
    
    // 检查是否已经初始化
    jlong currentHandle = env->GetLongField(thiz, handleField);
    if (currentHandle != 0) {
        // 先释放已有引擎
        ScoreEngine* oldEngine = reinterpret_cast<ScoreEngine*>(currentHandle);
        release_score_engine(oldEngine);
        env->SetLongField(thiz, handleField, 0);
    }
    
    // 获取基准文件数据
    jbyte* baseFileData = env->GetByteArrayElements(jBaseFileData, nullptr);
    if (baseFileData == nullptr) {
        return JNI_FALSE;
    }
    
    // 转换Java配置为C++配置
    ScoreEngineConfig config = convertJavaConfig(env, jConfig);
    
    // 调用C++初始化函数，传入实际的数据大小
    ScoreEngine* engine = init_score_engine(baseFileData, jDataSize, jFrameSize, jSampleRate, &config);

//    CharEntry * char_entry_array = (CharEntry *)((char*)baseFileData + engine->data_offset);
//
//    int char_entry_count = (jDataSize - engine->data_offset) / sizeof(CharEntry);
//    for (int i = 0; i < char_entry_count; i++) {
//        CharEntry char_entry = char_entry_array[i];
//        LOGD("CharEntry %d: start_frame %d,end_frame %d,energy %d pitch %d", i, char_entry.start_frame, char_entry.end_frame,char_entry.energy,char_entry.pitch);
//    }
    // 释放数组资源
    env->ReleaseByteArrayElements(jBaseFileData, baseFileData, JNI_ABORT);
    
    // 释放动态分配的配置资源
    if (config.custom_weights != nullptr) {
        delete config.custom_weights;
    }
    if (config.custom_factors != nullptr) {
        delete config.custom_factors;
    }
    
    // 检查引擎是否成功初始化
    if (engine == nullptr) {
        return JNI_FALSE;
    }
    
    // 保存引擎句柄到Java对象
    env->SetLongField(thiz, handleField, reinterpret_cast<jlong>(engine));
    
    return JNI_TRUE;
}

/**
 * 获取初始化错误码
 */
JNI_METHOD(jint, getInitError)(JNIEnv* env, jobject thiz) {
    return get_init_error();
}

/**
 * 实时评分接口函数
 */
JNI_METHOD(jint, scoreAudioFrame)(JNIEnv* env, jobject thiz, jfloatArray jAudioFrame, 
                                 jint jFrameIndex, jobject jResult, jfloatArray jSpectrum) {
    // 获取Java类引用
    jclass clazz = env->GetObjectClass(thiz);
    
    // 获取引擎句柄字段ID
    jfieldID handleField = env->GetFieldID(clazz, "nativeEngineHandle", "J");
    
    // 获取引擎句柄
    jlong handle = env->GetLongField(thiz, handleField);
    if (handle == 0) {
        return ERROR_ENGINE_NOT_INITIALIZED;
    }
    
    ScoreEngine* engine = reinterpret_cast<ScoreEngine*>(handle);
    
    // 获取音频帧数据
    jsize frameSize = env->GetArrayLength(jAudioFrame);
    float* audioFrame = env->GetFloatArrayElements(jAudioFrame, nullptr);
    
    // 准备频谱数据缓冲区
    float* spectrum = nullptr;
    if (jSpectrum != nullptr) {
        spectrum = env->GetFloatArrayElements(jSpectrum, nullptr);
    }
    
    // 创建评分结果
    ScoreResult cResult;
    memset(&cResult, 0, sizeof(ScoreResult));
    
    // --- TIMING START ---
//    auto start_time = std::chrono::high_resolution_clock::now();
    
    // 调用C++评分函数
    int result = score_audio_frame(engine, audioFrame, jFrameIndex, &cResult, spectrum);

     // --- TIMING START ---
//    auto end_time = std::chrono::high_resolution_clock::now();
//    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
//    LOGD("score_audio_frame time: %lld ms.",duration.count()/1000);
    // --- TIMING END ---
    // 释放数组资源
    env->ReleaseFloatArrayElements(jAudioFrame, audioFrame, JNI_ABORT);
    
    // 填充Java结果对象
    if (result == ERROR_SUCCESS && jResult != nullptr) {
        fillJavaResult(env, cResult, jResult);
    }
    
    // 释放频谱数据缓冲区
    if (spectrum != nullptr) {
        env->ReleaseFloatArrayElements(jSpectrum, spectrum, 0);
    }
    
    return result;
}

/**
 * 释放评分引擎资源
 */
JNI_METHOD(void, releaseScoreEngine)(JNIEnv* env, jobject thiz) {
    // 获取Java类引用
    jclass clazz = env->GetObjectClass(thiz);
    
    // 获取引擎句柄字段ID
    jfieldID handleField = env->GetFieldID(clazz, "nativeEngineHandle", "J");
    
    // 获取引擎句柄
    jlong handle = env->GetLongField(thiz, handleField);
    if (handle != 0) {
        ScoreEngine* engine = reinterpret_cast<ScoreEngine*>(handle);
        release_score_engine(engine);
        env->SetLongField(thiz, handleField, 0);
    }
}

// 获取data_offset
JNI_METHOD(jint, getDataOffset)(JNIEnv* env, jobject thiz)
{
    jclass clazz = env->GetObjectClass(thiz);
    jfieldID handleField = env->GetFieldID(clazz, "nativeEngineHandle", "J");
    jlong handle = env->GetLongField(thiz, handleField);
    if (handle == 0) return 0;
    ScoreEngine* engine = reinterpret_cast<ScoreEngine*>(handle);
    if (!engine) return 0;
    return (jint)engine->data_offset;
}

/**
 * 获取引擎版本信息
 */
JNI_METHOD(jstring, getEngineVersion)(JNIEnv* env, jobject thiz) {
    return env->NewStringUTF(ENGINE_VERSION);
}